package com.yugao.fintech.framework.log.core.customizer;

import brave.TracingCustomizer;
import brave.baggage.*;
import brave.context.slf4j.MDCScopeDecorator;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.ThreadLocalCurrentTraceContext;
import com.yugao.fintech.framework.log.common.LogCons;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 扩展mdc上下文字段, 从请求头中获取 {@link LogCons#LOGIN_ID} 字段并添加到mdc上下文中
 *
 * 发现在网关中如果定义了traceId的BaggageField字段, traceId有时会失效, 所以这里单独定义
 */
@Configuration
public class MdcFieldCustomizerConfig {

    @Bean
    BaggageField loginIdField() {
        return BaggageField.create(LogCons.LOGIN_ID);
    }

    @Bean
    CurrentTraceContext.ScopeDecorator mdcScopeDecorator(BaggageField loginIdField, BaggageField traceIdField) {
        return MDCScopeDecorator.newBuilder()
                .clear()
                .add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(loginIdField)
                        .flushOnUpdate()
                        .build())
                // 必须加上traceId, 消息队列在消费的时候才会将traceContent中的traceIdString 写入到mdc上下文
                .add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(BaggageFields.TRACE_ID)
                        .build())
                .build();
    }

    @Bean
    public CorrelationScopeCustomizer correlationScopeCustomizer() {
        return builder ->
                builder.clear().add(CorrelationScopeConfig.SingleCorrelationField.create(BaggageFields.TRACE_ID))
                ;
    }

    /**
     * Sleuth 默认只将 traceId 和 spanId 添加进入了 MDC {@link brave.baggage.CorrelationScopeDecorator.Builder }
     * 我们需要将 loginId 输出到 logback 中，所以在这里自定义添加 loginId 到MDC中
     * Sleuth 为我们提供了自定义扩展接口 CorrelationScopeCustomizer，实现它即可对 builder 进行操作
     */
    @Bean
    public TracingCustomizer tracingCustomizer(BaggageField loginIdField,
                                               Propagation.Factory factory,
                                               CurrentTraceContext.ScopeDecorator mdcScopeDecorator) {
        return builder ->
                builder.propagationFactory(BaggagePropagation.newFactoryBuilder(factory)
                        .add(BaggagePropagationConfig.SingleBaggageField.remote(loginIdField))
                        .build())
                .currentTraceContext(ThreadLocalCurrentTraceContext.newBuilder()
                        .addScopeDecorator(mdcScopeDecorator)
                        .build());
    }


}
